import os
import fileinput

from neo4j import GraphDatabase

# uri = "neo4j://192.168.88.161:7687" # 4.x版本用这个
uri = "bolt://192.168.88.161:7687" # 3.x版本用这个
driver = GraphDatabase.driver(uri, auth=("neo4j", "123456"), max_connection_lifetime=1000)

def load_dict_from_csv(path:str)->dict:
    """
    加载疾病的csv文件中存储的正症状，生成字典，key就是疾病名称，value就是症状列表
    :param path: 所有审核过的csv存储的路径
    :return: 字典
    """
    # 1 获取疾病名称
    disease_csv_list = os.listdir(path)
    print(disease_csv_list[:10])
    disease_list = [item.split('.')[0] for item in disease_csv_list]
    print(disease_list[:10])

    # 2 循环读取每个csv（疾病）每一行
    symptom_list = []
    for item in disease_csv_list:
        symptom = list(map(lambda x: x.strip(),
            fileinput.FileInput(os.path.join(path, item))))
        # print(symptom)
        symptom = list(filter(lambda x: 0<len(x)<100, symptom))
        # print(symptom)
        symptom_list.append(symptom)

    # 3 转成字典
    return dict(zip(disease_list, symptom_list))


def write_dict_to_neo4j(disease_symptom_dict:dict)->None:
    """
    接收疾病和症状的字典，写入到neo4j
    :param disease_symptom_dict:
    :return:
    """
    with driver.session() as sess:
        for key, value in disease_symptom_dict.items():
            # key就是一个疾病名称，value是症状的列表
            cypher = "MERGE (a:Disease{name:%r}) RETURN a" % key
            sess.run(cypher)
            for v in value:
                # 创建症状节点
                cypher = "MERGE (b:Symptom{name:%r}) RETURN b" % v
                sess.run(cypher)

                # 创建节点的关系
                cypher = "MATCH (a:Disease{name:%r}) MATCH (b:Symptom{name:%r}) \
                          WITH a,b MERGE(a)-[r:dis_to_sym]-(b)" %(key, v)
                sess.run(cypher)

        cypher = "CREATE INDEX ON:Disease(name)"
        sess.run(cypher)
        cypher = "CREATE INDEX ON:Symptom(name)"
        sess.run(cypher)


if __name__ == '__main__':
    d = load_dict_from_csv('/root/ai16/doctor_offline/structured/reviewed')
    # print(d)
    write_dict_to_neo4j(d)